#!/usr/bin/env python3
"""
OpenAI-Python 2.6.0 流式验证脚本
"""
from openai import OpenAI
import sys

client = OpenAI(
    base_url="http://localhost:11434/v1",
    api_key="ollama"          # 随意填
)

# MODEL = "deepseek-coder"      # 确保已 ollama pull deepseek-coder
MODEL = "deepseek-r1:32b"      # 确保已 ollama pull deepseek-coder

def stream_chat(prompt: str):
    try:
        print("【Ollama→Flask】流式响应：")
        stream = client.chat.completions.create(
            model=MODEL,
            messages=[
                {
                    "role": "system",
                    "content": "You are a helpful assistant."
                },
                {
                    "role": "user",
                    "content": "Hello!"
                }
                ],
            stream=True
        )
        for chunk in stream:
            delta = chunk.choices[0].delta.content
            if delta:
                print(delta, end="", flush=True)
        print("\n【END】")
    except Exception as e:
        print("【ERROR】", e)

if __name__ == "__main__":
    prompt = " ".join(sys.argv[1:]) if len(sys.argv) > 1 else "用 Python 写个快速排序，只要代码。"
    stream_chat(prompt)